from werkzeug.utils import secure_filename  # 添加 secure_filename
import uuid  # 添加 uuid
import os   # 添加 os
from flask_wtf.file import FileField

from myapp.views.baseSQLA import MyappSQLAInterface as SQLAInterface
from flask_babel import gettext as __
from flask_babel import lazy_gettext as _
from myapp.models.model_job import Repository,Images
from myapp.views.view_team import Creator_Filter, Project_Join_Filter, Project_Filter
from myapp import app, appbuilder, db
from wtforms.validators import DataRequired, Length, Regexp
from wtforms import StringField, SelectField
import pysnooper
import json
from flask import redirect, flash, request
from myapp.utils.py.py_k8s import K8s
from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
from myapp.forms import MyBS3TextAreaFieldWidget, MySelect2Widget
from flask_appbuilder import expose
from flask_appbuilder.actions import action
from .baseApi import MyappModelRestApi
from flask import g
from .base import (
    DeleteMixin,
    MyappFilter,
    MyappModelView,
)

conf = app.config




class Repository_ModelView_Base():
    datamodel = SQLAInterface(Repository)

    label_title = _('仓库')
    base_permissions = ['can_add', 'can_edit', 'can_delete', 'can_list', 'can_show']
    base_order = ('id', 'desc')
    # base_filters = [["id", Creator_Filter, lambda: []]]
    order_columns = ['id']
    search_columns = ['name', 'server', 'hubsecret', 'user']
    list_columns = ['name', 'server', 'hubsecret_url', 'creator', 'modified']
    cols_width = {
        "name": {"type": "ellip2", "width": 150},
        "hubsecret": {"type": "ellip2", "width": 350},
        "hubsecret_url": {"type": "ellip2", "width": 150},
    }
    show_exclude_columns = ['password']
    add_columns = ['name', 'server', 'user', 'password', 'hubsecret']
    edit_columns = add_columns

    spec_label_columns={
        "server": _('仓库'),
        "user": _("用户名"),
        "hubsecret": 'k8s hubsecret',
        "hubsecret_url": 'k8s hubsecret'
    }

    add_form_extra_fields = {
        "server": SelectField(
            _('服务地址'),
            widget=MySelect2Widget(can_input=True),
            default='harbor.oa.com/cube-studio/',
            choices=[['harbor.oa.com/cube-studio/','harbor.oa.com/cube-studio/'],['ccr.ccs.tencentyun.com/cube-studio/','ccr.ccs.tencentyun.com/cube-studio/'],['registry.docker-cn.com','registry.docker-cn.com']],
            # description= _("镜像仓库地址")
            description= _("镜像仓库地址，示例：")+conf.get('REPOSITORY_ORG','')
        ),
        "user": StringField(
            _('用户名'),
            default='',
            widget=BS3TextFieldWidget(),
            description= _("镜像仓库的用户名")
        ),
        "password": StringField(
            _('密码'),
            default='',
            widget=BS3TextFieldWidget(),
            description= _("镜像仓库的链接密码")
        )
    }

    edit_form_extra_fields = add_form_extra_fields

    # @pysnooper.snoop()
    def set_column(self):
        self.add_form_extra_fields['name'] = StringField(
            _('名称'),
            default=g.user.username + "-",
            widget=BS3TextFieldWidget(),
            description= _("仓库名称")
        )

        self.add_form_extra_fields['hubsecret'] = StringField(
            "hubsecret",
            default=g.user.username + "-hubsecret",
            widget=BS3TextFieldWidget(),
            description= _("在k8s中创建的hub secret"),
            validators=[Regexp("^[a-z][a-z0-9\-]*[a-z0-9]$"), Length(1, 54), DataRequired()]
        )

    pre_add_web = set_column

    def check_edit_permission(self, item):
        if not g.user.is_admin() and g.user.username != item.created_by.username:
            return False
        return True
    check_delete_permission = check_edit_permission

    # create hubsecret
    # @pysnooper.snoop()
    def apply_hubsecret(self, repo):
        from myapp.utils.py.py_k8s import K8s
        all_cluster = conf.get('CLUSTERS', {})
        all_kubeconfig = [all_cluster[cluster].get('KUBECONFIG', '') for cluster in all_cluster] + ['']
        all_kubeconfig = list(set(all_kubeconfig))
        for kubeconfig in all_kubeconfig:
            try:
                k8s = K8s(kubeconfig)
                namespaces = conf.get('HUBSECRET_NAMESPACE',[])
                for namespace in namespaces:
                    try:
                        server = repo.server[:repo.server.index('/')] if '/' in repo.server else repo.server
                        # print(server)
                        k8s.apply_hubsecret(namespace=namespace,
                                            name=repo.hubsecret,
                                            user=repo.user,
                                            password=repo.password,
                                            server=server
                                            )
                    except Exception as e1:
                        print(e1)
            except Exception as e:
                print(e)

    def post_add(self, item):
        self.apply_hubsecret(item)

    def post_update(self, item):
        self.apply_hubsecret(item)


class Repository_ModelView_Api(Repository_ModelView_Base, MyappModelRestApi):
    datamodel = SQLAInterface(Repository)

    route_base = '/repository_modelview/api'

appbuilder.add_api(Repository_ModelView_Api)



class Images_ModelView_Base():
    label_title = _('镜像')
    datamodel = SQLAInterface(Images)

    list_columns = ['project','images_url', 'upload_type', 'upload_status', 'upload_progress', 'creator', 'modified']
    cols_width = {
        "images_url": {"type": "ellip2", "width": 500},
    }
    search_columns = ['created_by', 'project', 'repository', 'name', 'describe']
    base_order = ('id', 'desc')
    order_columns = ['id']
    add_columns = ['project','repository', 'name', 'describe', 'dockerfile', 'gitpath', 'upload_type', 'local_path', 'upload_status']
    edit_columns = ['project','repository', 'name', 'describe', 'dockerfile', 'gitpath', 'upload_type', 'local_path', 'upload_status', 'upload_progress']
    spec_label_columns={
        "project": _("功能分类"),
        "upload_type": _("上传类型"),
        "local_path": _("本地路径"),
        "file_size": _("文件大小"),
        "upload_status": _("上传状态"),
        "upload_progress": _("上传进度"),
        "name": _("镜像名称"),
        "describe": _("描述"),
        "dockerfile": _("Dockerfile"),
        "gitpath": _("Git路径"),
        "repository": _("镜像仓库"),
        "entrypoint": _("启动命令")
    }
    add_form_query_rel_fields = {
        "project": [["name", Project_Filter, 'job-template']]
    }
    edit_form_query_rel_fields = add_form_query_rel_fields


    add_form_extra_fields = {
        "upload_type": SelectField(
            _('创建方式'),
            choices=[('dockerfile', 'Dockerfile构建'), ('upload', '本地上传')],
            default='dockerfile',
            widget=MySelect2Widget(),
            description=_('选择镜像创建方式')
        ),
        "dockerfile": StringField(
            'dockerfile',
            description= _('镜像的构建Dockerfile全部内容，/mnt/$username/是个人存储目录，可以从此目录下复制文件到镜像中'),
            default='',
            widget=MyBS3TextAreaFieldWidget(rows=10),
            render_kw={"class": "dockerfile-fields"}
        ),
        "gitpath": StringField(
            _('Git路径'),
            default='',
            widget=BS3TextFieldWidget(),
            description=_('Git仓库路径'),
            render_kw={"class": "dockerfile-fields"}
        ),
        "name": StringField(
            _('名称'),
            description= _('镜像名称全称，例如ubuntu:20.04'),
            default='',
            widget=BS3TextFieldWidget(),
        ),
        "entrypoint": StringField(
            _('启动命令'),
            description= _('镜像的入口命令，直接写成单行字符串，例如python xx.py，无需添加[]'),
            default='',
            widget=BS3TextFieldWidget(),
        ),
        # 添加文件上传字段
        "image_file": FileField(
            _('镜像文件'),
            description=_('选择本地镜像文件（.tar、.tar.gz、.zip格式）'),
            render_kw={"class": "upload-fields", "style": "display:none;"}
        ),
        "local_path": StringField(
            _('本地路径'),
            default='',
            widget=BS3TextFieldWidget(),
            description=_('本地镜像文件路径（上传时自动填充）'),
            render_kw={"style": "display:none;"}
        ),
        "upload_status": SelectField(
            _('上传状态'),
            choices=[('pending', '等待处理'), ('uploaded', '文件已上传'), ('processing', '正在处理'), ('success', '处理完成'), ('failed', '处理失败')],
            default='pending',
            widget=MySelect2Widget(),
            description=_('镜像上传处理状态')
        )
    }

    edit_form_extra_fields = add_form_extra_fields
    # base_filters = [["id", Creator_Filter, lambda: []]]

    def check_edit_permission(self, item):
        if not g.user.is_admin() and g.user.username != item.created_by.username:
            return False
        return True
    check_delete_permission = check_edit_permission




class Images_ModelView_Api(Images_ModelView_Base, MyappModelRestApi):
    datamodel = SQLAInterface(Images)
    route_base = '/images_modelview/api'
    
    def _import_image_to_registry(self, image):
        """将上传的镜像导入到Docker仓库"""
        try:
            import subprocess
            import tempfile
            import shutil
            
            # 解压镜像文件
            with tempfile.TemporaryDirectory() as temp_dir:
                if image.local_path.endswith('.zip'):
                    import zipfile
                    with zipfile.ZipFile(image.local_path, 'r') as zip_ref:
                        zip_ref.extractall(temp_dir)
                elif image.local_path.endswith(('.tar', '.tar.gz', '.tgz')):
                    import tarfile
                    with tarfile.open(image.local_path, 'r') as tar_ref:
                        tar_ref.extractall(temp_dir)
                
                # 查找镜像文件
                image_files = []
                for root, dirs, files in os.walk(temp_dir):
                    for file in files:
                        if file.endswith('.tar'):
                            image_files.append(os.path.join(root, file))
                
                if not image_files:
                    raise Exception("未找到有效的镜像文件")
                
                # 导入镜像到Docker
                for image_file in image_files:
                    cmd = f"docker load -i {image_file}"
                    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
                    if result.returncode != 0:
                        raise Exception(f"导入镜像失败：{result.stderr}")
                
                # 标记镜像并推送到仓库
                repository = image.repository
                if repository:
                    # 重新标记镜像
                    new_tag = f"{repository.server}{image.name}"
                    cmd = f"docker tag {image.name} {new_tag}"
                    subprocess.run(cmd, shell=True)
                    
                    # 推送到仓库
                    cmd = f"docker push {new_tag}"
                    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
                    if result.returncode != 0:
                        raise Exception(f"推送镜像失败：{result.stderr}")
            
            return True
        except Exception as e:
            print(f"导入镜像到仓库失败：{str(e)}")
            return False
    
    def post_add(self, item):
        """处理添加后的逻辑"""
        if item.upload_type == 'upload' and hasattr(request, 'files') and 'image_file' in request.files:
            file = request.files['image_file']
            if file and file.filename:
                # 处理文件上传
                upload_dir = os.path.join(conf.get('UPLOAD_FOLDER', '/tmp'), 'images', g.user.username)
                os.makedirs(upload_dir, exist_ok=True)
                
                filename = secure_filename(file.filename)
                unique_filename = f"{uuid.uuid4().hex}_{filename}"
                file_path = os.path.join(upload_dir, unique_filename)
                file.save(file_path)
                
                # 更新记录
                item.local_path = file_path
                item.file_size = os.path.getsize(file_path)
                item.upload_status = 'uploaded'
                db.session.commit()
                
                # 启动后台处理
                self._process_uploaded_image(item.id)
        
        return super().post_add(item)
    
    @expose("/upload_image/", methods=["POST"])
    def upload_image(self):
        """上传本地镜像文件"""
        try:
            # 检查上传的文件
            if 'image_file' not in request.files:
                return self.response(400, message="没有选择文件")
            
            file = request.files['image_file']
            if file.filename == '':
                return self.response(400, message="没有选择文件")
            
            # 获取其他参数
            image_name = request.form.get('image_name', '')
            image_describe = request.form.get('image_describe', '')
            project_id = request.form.get('project_id')
            repository_id = request.form.get('repository_id')
            
            if not all([image_name, project_id, repository_id]):
                return self.response(400, message="缺少必要参数")
            
            # 验证文件类型（支持.zip文件）
            allowed_extensions = {'.tar', '.tar.gz', '.tgz', '.zip'}
            file_ext = os.path.splitext(file.filename)[1].lower()
            if file_ext not in allowed_extensions:
                return self.response(400, message="不支持的文件格式，请上传.tar、.tar.gz或.zip文件")
            
            # 创建上传目录
            upload_dir = os.path.join(conf.get('UPLOAD_FOLDER', '/tmp'), 'images', g.user.username)
            os.makedirs(upload_dir, exist_ok=True)
            
            # 保存文件
            filename = secure_filename(file.filename)
            unique_filename = f"{uuid.uuid4().hex}_{filename}"
            file_path = os.path.join(upload_dir, unique_filename)
            file.save(file_path)
            
            # 获取文件大小
            file_size = os.path.getsize(file_path)
            
            # 创建数据库记录
            image_record = Images(
                name=image_name,
                describe=image_describe,
                project_id=int(project_id),
                repository_id=int(repository_id),
                upload_type='upload',
                local_path=file_path,
                file_size=file_size,
                upload_status='uploaded',
                upload_progress=0,
                created_by_fk=g.user.id
            )
            
            db.session.add(image_record)
            db.session.commit()
            
            # 启动后台处理任务
            self._process_uploaded_image(image_record.id)
            
            return self.response(200, message="文件上传成功，正在处理中", data={"image_id": image_record.id})
            
        except Exception as e:
            db.session.rollback()
            return self.response(500, message=f"上传失败：{str(e)}")
    
    def _process_uploaded_image(self, image_id):
        """处理上传的镜像文件（后台任务）"""
        try:
            image = db.session.query(Images).get(image_id)
            if not image:
                return
            
            # 更新状态为处理中
            image.upload_status = 'processing'
            image.upload_progress = 10
            db.session.commit()
            
            # 这里添加实际的镜像处理逻辑
            # 例如：解压文件、导入到Docker仓库等
            
            # 模拟处理过程
            import time
            time.sleep(2)  # 模拟处理时间
            
            # 更新进度
            image.upload_progress = 50
            db.session.commit()
            
            # 实际的镜像导入逻辑应该在这里
            # self._import_image_to_registry(image)
            
            # 完成处理
            image.upload_status = 'success'
            image.upload_progress = 100
            db.session.commit()
            
        except Exception as e:
            # 处理失败
            image.upload_status = 'failed'
            image.upload_progress = 0
            db.session.commit()
            print(f"处理镜像失败：{str(e)}")
    
    @expose("/upload_progress/<int:image_id>", methods=["GET"])
    def upload_progress(self, image_id):
        """获取上传进度"""
        try:
            image = db.session.query(Images).get(image_id)
            if not image:
                return self.response(404, message="镜像记录不存在")
            
            return self.response(200, data={
                "progress": image.upload_progress,
                "status": image.upload_status,
                "message": self._get_status_message(image.upload_status)
            })
        except Exception as e:
            return self.response(500, message=f"获取进度失败：{str(e)}")
    
    def _get_status_message(self, status):
        """获取状态描述"""
        status_messages = {
            'pending': '等待处理',
            'uploaded': '文件已上传',
            'processing': '正在处理镜像',
            'success': '处理完成',
            'failed': '处理失败'
        }
        return status_messages.get(status, '未知状态')

class Images_ModelView(Images_ModelView_Base, MyappModelView):
    datamodel = SQLAInterface(Images)
    
    def _get_add_widget(self, **args):
        """重写添加页面，添加动态字段控制"""
        widgets = super()._get_add_widget(**args)
        
        upload_script = """
        <script>
        $(document).ready(function() {
            // 监听创建方式选择变化
            $('#upload_type').on('change', function() {
                var selectedType = $(this).val();
                
                if (selectedType === 'dockerfile') {
                    $('.dockerfile-fields').show();
                    $('.upload-fields').hide();
                    $('#dockerfile').prop('required', true);
                    $('#image_file').prop('required', false);
                } else if (selectedType === 'upload') {
                    $('.dockerfile-fields').hide();
                    $('.upload-fields').show();
                    $('#dockerfile').prop('required', false);
                    $('#image_file').prop('required', true);
                }
            });
            
            // 初始化显示状态
            $('#upload_type').trigger('change');
            
            // 处理文件上传
            $('#image_file').on('change', function() {
                var file = this.files[0];
                if (file) {
                    var allowedTypes = ['.tar', '.tar.gz', '.tgz', '.zip'];
                    var fileName = file.name.toLowerCase();
                    var isValidType = allowedTypes.some(type => fileName.endsWith(type));
                    
                    if (!isValidType) {
                        alert('请选择有效的镜像文件格式（.tar、.tar.gz、.zip）');
                        $(this).val('');
                        return;
                    }
                    
                    $('#local_path').val(file.name);
                }
            });
        });
        </script>
        """
        
        widgets['upload_script'] = upload_script
        return widgets
    
    # # 添加上传按钮到操作列表
    # @action("upload_local_image", "上传本地镜像", "上传本地Docker镜像文件", "fa-upload", multiple=False)
    # def upload_local_image(self, items):
    #     """上传本地镜像的操作按钮"""
    #     return redirect('/images_modelview/list/')
    
      
    
# 注册Web视图
appbuilder.add_view(
    Images_ModelView,
    "镜像管理",
    icon="fa-image",
    category="镜像",
    category_icon="fa-image"
)

appbuilder.add_api(Images_ModelView_Api)



